#include <array>
#include <shared_mutex>
#include <unordered_map>

template <typename Key, typename Value>
class StripedLockingKVDB : public KVDatabase<Key, Value> {
private:
    static constexpr size_t kNumShards = 64;  // 分片数（建议为CPU核心数2倍）
    
    struct Shard {
        std::unordered_map<Key, Value> data;
        mutable std::shared_mutex mutex;
    };
    
    std::array<Shard, kNumShards> shards_;
    
    // 根据Key计算分片索引
    size_t getShardIndex(const Key& key) const {
        return std::hash<Key>{}(key) % kNumShards;
    }

public:
    StripedLockingKVDB() = default;

    // 插入或更新键值对
    bool put(const Key& key, const Value& value) override {
        auto& shard = shards_[getShardIndex(key)];
        std::unique_lock lock(shard.mutex);
        shard.data[key] = value;
        return true;
    }

    // 查询键值对
    bool get(const Key& key, Value& value) const override {
        auto& shard = shards_[getShardIndex(key)];
        std::shared_lock lock(shard.mutex);
        auto it = shard.data.find(key);
        if (it != shard.data.end()) {
            value = it->second;
            return true;
        }
        return false;
    }

    // 判断键是否存在
    bool exists(const Key& key) const override {
        auto& shard = shards_[getShardIndex(key)];
        std::shared_lock lock(shard.mutex);
        return shard.data.count(key) > 0;
    }

    // 删除键值对
    bool remove(const Key& key) override {
        auto& shard = shards_[getShardIndex(key)];
        std::unique_lock lock(shard.mutex);
        return shard.data.erase(key) > 0;
    }

    // 序列化（需要锁定所有分片）
    bool serialize(std::string& data_str) const override {
        std::vector<std::unique_lock<std::shared_mutex>> locks;
        for (auto& shard : shards_) {
            locks.emplace_back(shard.mutex);
        }
        
        try {
            std::ostringstream oss;
            boost::archive::text_oarchive oa(oss);
            for (const auto& shard : shards_) {
                oa << shard.data;
            }
            data_str = oss.str();
            return true;
        } catch (...) {
            return false;
        }
    }

    // 反序列化（需要锁定所有分片）
    bool deserialize(const std::string& data_str) override {
        std::vector<std::unique_lock<std::shared_mutex>> locks;
        for (auto& shard : shards_) {
            locks.emplace_back(shard.mutex);
        }

        try {
            std::istringstream iss(data_str);
            boost::archive::text_iarchive ia(iss);
            for (auto& shard : shards_) {
                ia >> shard.data;
            }
            return true;
        } catch (...) {
            return false;
        }
    }

    // 获取总键值对数
    size_t size() const override {
        size_t total = 0;
        for (auto& shard : shards_) {
            std::shared_lock lock(shard.mutex);
            total += shard.data.size();
        }
        return total;
    }
};